#include "dyn_roots_multiDim_sta.h"

MultiSTA::MultiSTA(MultiMaps& map,
                   colvec init,
                   const double tol,
                   const double para,
                   const int p,
                   const int Nmax)
    : multiMap(map)
{
    setm_InitialValue(init);
    setTolerance(tol);
    setLambda(para);
    setMaxIterations(Nmax);
    setPeriod(p);
    prepareReflectionMatrices(m_initialValue.n_elem);
}

int MultiSTA::solve()
{
    // We start the clock
    start = clock();

    // We initialize the placeholders for our
    // return values and the return code.
    int returnCode(1);


    // Since we only implement the 2D version,
    // there are some literal constants in the algorithm.
    // We first resize the fixedPoints result matrix to fit
    // the algorithm and fill it with the initial values.
    fixedPoints.set_size(5,2);

    for (int k=0;k<5;k++)
    {
        if (m_initialValue.is_colvec())
        {
            rowvec init = m_initialValue.t();
            fixedPoints.row(k) = init;
        }
        else
        {
            rowvec init = m_initialValue;
            fixedPoints.row(k) = init;
        }
    }

    for (unsigned int i=0;i<reflectionMatrices.n_slices;i++)
    {
        for (int j=0;j<maxIterations;j++)
        {
            colvec dx = iterate(fixedPoints.row(i).t(),i);
            iterations++;

            fixedPoints.row(i) += dx.t();

            // We stop if we attain convergence or diverge.
            if (norm(dx,2) < tolerance)
            {
                returnCode = 0;
                break;
            }

            if (norm(dx,2) > 100)
            {
                fixedPoints.row(i).zeros();
                break;
            }
        }
    }

    // We stop the clock
    end = clock();
    processorTime = ((double) (end-start)) / CLOCKS_PER_SEC;
    return returnCode;
}

colvec MultiSTA::iterate(const colvec X, int sliceIndex)
{
    colvec iterationValue = multiMap.iterate(X,period)-X;
    colvec dx = lambda*reflectionMatrices.slice(sliceIndex)*iterationValue;

    return dx;
}

int MultiSTA::prepareReflectionMatrices(int dim)
{
    int returnCode(1);

    switch (dim)
    {
    case 1:
    {
        mat reflection1(1,1), reflection2(1,1);
        reflection1(0,0) = 1;
        reflection2(0,0) = -1;

        reflectionMatrices.slice(0) = reflection1;
        reflectionMatrices.slice(1) = reflection2;
        break;
    }
    case 2:
    {
        reflectionMatrices.set_size(2,2,5);
        mat reflection1(2,2), reflection2(2,2), reflection3(2,2), reflection4(2,2), reflection5(2,2);
        reflection1 = eye(2,2);

        reflection2 << -1.0 << 0.0 << endr
                    << 0.0  << 1.0 << endr;

        reflection3 << 1.0 << 0.0 << endr
                    << 0.0 << -1.0 << endr;

        reflection4 << 0.0 << -1.0 << endr
                    << -1.0 << 0.0 << endr;

        reflection5 << 0.0 << 1.0 << endr
                    << 1.0 << 0.0 << endr;

        reflectionMatrices.slice(0) = reflection1;
        reflectionMatrices.slice(1) = reflection2;
        reflectionMatrices.slice(2) = reflection3;
        reflectionMatrices.slice(3) = reflection4;
        reflectionMatrices.slice(4) = reflection5;
        returnCode = 0;
        break;
    }
    default:
    {
        cout << "Not supported yet." << endl;
        break;
    }
    }
    return returnCode;
}
